
#include <ntifs.h>

#include "common.h"
#include "winver.h"
#include "helper.h"
#include "object.h"


//PULONG ObTypeIndexTable = 0x83d598c0;

POBJECT_DIRECTORY ObpRootDirectoryObject = 0x8d401120;

#define WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_OFFSET    0x60
#define WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_DumpProcedure_OFFSET 0x2c

//
// for win7
//

#define WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_OFFSET_WIN7    0x28
#define WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_DumpProcedure_OFFSET_WIN7 0x30

NTSTATUS 
QueryObjectsInfo(
    OUT PVOID ObjectInfoBuffer,
    IN ULONG BufferLength,
    OUT PULONG Length      
)
{
	POBJECT_DIRECTORY pDirectory = NULL;
	POBJECT_DIRECTORY_ENTRY pEntry = NULL, DescentDirectory = NULL;
	POBJECT_HEADER_NAME_INFO pName;
	POBJECT_HEADER pObjHeader;
	PWCHAR szName;
	int Depth = 1, Bucket, DepthIndex;
    ULONG TotalLength = 0;
    POBJECT_INFO pObjectInfo;
    NTSTATUS ntStatus = STATUS_SUCCESS;
    
    PKDDEBUGGER_DATA64 KdDebugData = GetKdDebugerData();
    if (!KdDebugData || !KdDebugData->ObpRootDirectoryObject){
        return STATUS_UNSUCCESSFUL;
    }

    pDirectory = (POBJECT_DIRECTORY)KdDebugData->ObpRootDirectoryObject;

    //
    //  check parameter
    //
    
    if (!ObjectInfoBuffer || BufferLength < sizeof(OBJECT_ALL_INFO)){
        return STATUS_INVALID_PARAMETER;
    }
    
    //
    // add the length of the counts
    // the infos size see blow 
    //
    
    TotalLength += FIELD_OFFSET(OBJECT_ALL_INFO, Infos);

    __try{

	    while (pDirectory){
restart_dir_walk:

		    for (Bucket = 0; Bucket < NUMBER_HASH_BUCKETS; Bucket++) {
			    pEntry = pDirectory->HashBuckets[Bucket];
			    while (pEntry){
				    pObjHeader = OBJECT_TO_OBJECT_HEADER(pEntry->Object);
				    if (DescentDirectory) {
					    if (DescentDirectory == pEntry->Object){
						    DescentDirectory = NULL;
					    }
					    pEntry = pEntry->ChainLink;
					    continue;
				    }
				
				    //
				    //  for OBJECT_TYPE
				    //
				
				    if (pObjHeader->TypeIndex == OBJ_TYPE_TYPE){

					    for(DepthIndex = 1; DepthIndex < Depth; DepthIndex++) {
						    DbgPrint("  ");
                        }
					
                        pObjectInfo = (POBJECT_INFO)((ULONG)ObjectInfoBuffer + TotalLength);
                    
                        TotalLength += sizeof(OBJECT_INFO);

                        //
					    //  get name 
					    //
					
					    pName = (POBJECT_HEADER_NAME_INFO)ObQueryNameInfo(pEntry->Object);

					    //
					    //  print object type
					    //

					    if (pName){
                            TotalLength += pName->Name.Length + sizeof(WCHAR);
					    }

                        if (TotalLength > BufferLength){
                            ntStatus = STATUS_INFO_LENGTH_MISMATCH;
                            if (!ARGUMENT_PRESENT(Length)){
                                leave;
                            }
                        }else{

                            ULONG TypeInfo;
                            ULONG ProcStart;
                            PWCHAR pNameTmp = (PWCHAR)((ULONG)pObjectInfo + sizeof(OBJECT_INFO));
                            
                            //
                            //  win7 modified the OBJECT_TYPE structure
                            //  like is:
                            //  windows xp sp2
                            //  kd> dt _OBJECT_TYPE
                            //    nt!_OBJECT_TYPE
                            //       +0x000 Mutex            : _ERESOURCE
                            //       +0x038 TypeList         : _LIST_ENTRY
                            //       +0x040 Name             : _UNICODE_STRING
                            //       +0x048 DefaultObject    : Ptr32 Void
                            //       +0x04c Index            : Uint4B
                            //       +0x050 TotalNumberOfObjects : Uint4B
                            //       +0x054 TotalNumberOfHandles : Uint4B
                            //       +0x058 HighWaterNumberOfObjects : Uint4B
                            //       +0x05c HighWaterNumberOfHandles : Uint4B
                            //       +0x060 TypeInfo         : _OBJECT_TYPE_INITIALIZER
                            //       +0x0ac Key              : Uint4B
                            //       +0x0b0 ObjectLocks      : [4] _ERESOURCE
                            //  windows 7
                            //     nt!_OBJECT_TYPE
                            //       +0x000 TypeList         : _LIST_ENTRY
                            //       +0x008 Name             : _UNICODE_STRING
                            //       +0x010 DefaultObject    : Ptr32 Void
                            //       +0x014 Index            : UChar
                            //       +0x018 TotalNumberOfObjects : Uint4B
                            //       +0x01c TotalNumberOfHandles : Uint4B
                            //       +0x020 HighWaterNumberOfObjects : Uint4B
                            //       +0x024 HighWaterNumberOfHandles : Uint4B
                            //       +0x028 TypeInfo         : _OBJECT_TYPE_INITIALIZER
                            //       +0x078 TypeLock         : _EX_PUSH_LOCK
                            //       +0x07c Key              : Uint4B
                            //       +0x080 CallbackList     : _LIST_ENTRY
                            //
                            
                            if (IsWindowsXpTo2003()){
                                TypeInfo = *(PULONG)((ULONG)pEntry->Object + 
                                             WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_OFFSET);
                                ProcStart = TypeInfo + 
                                    WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_DumpProcedure_OFFSET;
                            }else{
                                TypeInfo = *(PULONG)((ULONG)pEntry->Object + 
                                             WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_OFFSET_WIN7);
                                ProcStart = TypeInfo + 
                                    WINDOWS_NT_KERNEL_OBJECT_TYPE_INITIALIZER_DumpProcedure_OFFSET_WIN7;
                            }

                            //
                            //  copy dispatch procedure
                            //

                            if (TypeInfo && ProcStart){
                                RtlCopyMemory(pObjectInfo->DiapathProc, ProcStart, sizeof(pObjectInfo->DiapathProc));
                            }

                            //
                            //  copy object type name
                            //
                        
                            if (pName){
                                pObjectInfo->ObjectTypeName.Length = pName->Name.Length + sizeof(WCHAR);
                                pObjectInfo->ObjectTypeName.MaximumLength = pObjectInfo->ObjectTypeName.Length;
                                RtlZeroMemory(pNameTmp, pName->Name.Length + sizeof(WCHAR));
                                RtlCopyMemory(pNameTmp, pName->Name.Buffer, pName->Name.Length);
                                pObjectInfo->ObjectTypeName.Buffer = pNameTmp;
                                pObjectInfo->Length = sizeof(*pObjectInfo) + pObjectInfo->ObjectTypeName.Length;
                            }else{
                                RtlZeroMemory(&pObjectInfo->ObjectTypeName, sizeof(UNICODE_STRING));
                                pObjectInfo->Length = sizeof(*pObjectInfo);
                            }
                        }
                        
                        //
                        //  next 
                        //
                        
					    pEntry = pEntry->ChainLink;
					    continue;

				    }else if (pObjHeader->TypeIndex == OBJ_TYPE_DIRECTORY){

					    for(DepthIndex =1; DepthIndex< Depth; DepthIndex++) {
						    DbgPrint("  ");
					    }

					    //
					    //  get name information 
					    //
					
					    //pName = (POBJECT_HEADER_NAME_INFO)ObQueryNameInfo(pEntry->Object);

					    //
					    //  print object type
					    //

					    //if (pName){
						//    szName = ExAllocatePoolWithTag(NonPagedPool, pName->Name.Length + sizeof(WCHAR),
                        //             0x80 );
						//    RtlZeroMemory(szName, pName->Name.Length + sizeof(WCHAR));
						//    RtlCopyMemory(szName, pName->Name.Buffer, pName->Name.Length);
						//    DbgPrint("Object Directory: name %ws\n", szName);
						//    ExFreePool(szName);
					    //}

					    //
					    //  next
					    //
					
					    pDirectory = pEntry->Object;
                        Depth++;
                        goto restart_dir_walk;
				
				    }else{
					    pEntry = pEntry->ChainLink;
				    }
			    }
		    }
		
		    //
		    //  next
		    //
		
		    Depth--;
		    pName = (POBJECT_HEADER_NAME_INFO)ObQueryNameInfo(pDirectory);
            DescentDirectory = pDirectory;
		    pDirectory = pName->Directory;
	    }

        //
        //  continue calculate the total length for output for user mode.
        //
                        
        if (ARGUMENT_PRESENT(Length)){
            *Length = TotalLength;
        }

    }except(EXCEPTION_EXECUTE_HANDLER){
        ntStatus = GetExceptionCode();
    }

    return ntStatus;
}